package com.gitee.wsl.ext.collection

import arrow.core.Either
import arrow.core.Option
import arrow.core.Some
import arrow.core.None
import arrow.core.raise.either
import arrow.core.raise.result

/**
 * Returns an Either of a list of B results of applying the given transform function
 * to each element(A) in the original collection.
 */
inline fun <E, A, B> Iterable<A>.traverse(f: (A) -> Either<E, B>): Either<E, List<B>> =
    let { l -> either { l.map { f(it).bind() } } }

/**
 * Returns a Result of a list of B results of applying the given transform function
 * to each element(A) in the original collection.
 */
inline fun <A, B> Iterable<A>.traverse(f: (A) -> Result<B>): Result<List<B>> =
    let { l -> result { l.map { f(it).bind() } } }

/**
 * Synonym for traverse((A)-> Either<E, B>): Either<E, List<B>>
 */
inline fun <E, A, B> Iterable<A>.traverseEither(f: (A) -> Either<E, B>): Either<E, List<B>> =
    traverse(f)


/**
 * Synonym for traverse((A)-> Result<B>): Result<List<B>>
 */
inline fun <A, B> Iterable<A>.traverseResult(f: (A) -> Result<B>): Result<List<B>> =
    traverse(f)

/**
 * Returns an Option of a list of B results of applying the given transform function
 * to each element(A) in the original collection.
 */

inline fun <A, B> Iterable<A>.traverse(f: (A) -> Option<B>): Option<List<B>> {
    val result = ArrayList<B>()
    for (element in this) {
        val mapped = f(element)
        if (mapped is Some) result.add(mapped.value)
        else return None
    }
    return Some(result)
}

/**
 * Synonym for traverse((A)-> Option<B>): Option<List<B>>
 */
inline fun <A, B> Iterable<A>.traverseOption(f: (A) -> Option<B>): Option<List<B>> =
    traverse(f)


/**
 * Map a function that returns an Either across the Sequence.
 */
@Suppress("EXTENSION_SHADOWED_BY_MEMBER")
inline fun <A, E, B> Sequence<A>.traverse(f: (A) -> Either<E, B>): Either<E, List<B>> {
    // Note: Using a mutable list here avoids the stackoverflows one can accidentally create when using
    //  Sequence.plus instead. But we don't convert the sequence to a list beforehand to avoid
    //  forcing too much of the sequence to be evaluated.
    val result = mutableListOf<B>()
    forEach { a ->
        when (val mapped = f(a)) {
            is Either.Right -> result.add(mapped.value)
            is Either.Left -> return@traverse mapped
        }
    }
    return Either.Right(result.toList())
}

/**
 * Synonym for traverse((A)-> Either<E, B>): Either<E, List<B>>
 */
@Suppress("EXTENSION_SHADOWED_BY_MEMBER")
inline fun <A, E, B> Sequence<A>.traverseEither(f: (A) -> Either<E, B>): Either<E, List<B>> =
    traverse(f)

/**
 * Map a function that returns an Option across the Sequence.
 */
inline fun <A, B> Sequence<A>.traverse(f: (A) -> Option<B>): Option<List<B>> {
    // Note: Using a mutable list here avoids the stackoverflows one can accidentally create when using
    //  Sequence.plus instead. But we don't convert the sequence to a list beforehand to avoid
    //  forcing too much of the sequence to be evaluated.
    val result = mutableListOf<B>()
    forEach { a ->
        when (val mapped = f(a)) {
            is Some -> result.add(mapped.value)
            is None -> return@traverse None
        }
    }
    return Some(result.toList())
}

/**
 * Synonym for traverse((A)-> Option<B>): Option<List<B>>
 */
@Suppress("EXTENSION_SHADOWED_BY_MEMBER")
inline fun <A, B> Sequence<A>.traverseOption(f: (A) -> Option<B>): Option<List<B>> =
    traverse(f)

